
# -------------------------- 1. 导入所有依赖库 --------------------------
import torch
import torch.nn.functional as F
from torchvision import transforms
from torchvision import datasets
from torch.utils.data import DataLoader
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np
# -------------------------- 2. 全局配置（图像预处理/绘图环境） --------------------------
# 1. 图像预处理配置：训练/测试用的归一化 + 显示用的反归一化
# 训练时归一化（MNIST数据集统计的均值和标准差）
train_transform = transforms.Compose([
    transforms.ToTensor(),  # 把PIL图像(0-255)转Tensor(0-1)
    transforms.Normalize(mean=(0.1307,), std=(0.3081,))  # 归一化到均值0、标准差1附近
])
# 显示图像时反归一化（还原为原始像素范围，避免显示偏色）
inv_transform = transforms.Compose([
    transforms.Normalize(mean=[0.0], std=[1/0.3081]),  # 先乘标准差：(x*std)
    transforms.Normalize(mean=[-0.1307], std=[1.0])    # 再加均值：(x*std + mean)，还原为0-1范围
])

# 2. 绘图环境初始化
plt.ion()  # 开启matplotlib交互模式，允许实时刷新图表
# 创建两个窗口：一个显示准确率曲线，一个显示图像及识别结果
fig_acc, ax_acc = plt.subplots(figsize=(8, 5))  # 准确率曲线窗口（8x5英寸）
fig_img, axs_img = plt.subplots(2, 5, figsize=(12, 6))  # 图像显示窗口（2行5列，共10个样本）
axs_img = axs_img.flatten()  # 把2D的坐标轴数组转1D，方便循环赋值

# 3. 训练基础参数
batch_size = 64  # 每次训练/测试的样本数量
epochs = 10  # 训练轮次（完整遍历训练集10次）
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # 自动选择GPU/CPU
print(f"Using device: {device}")  # 打印当前使用的设备


# -------------------------- 3. 数据加载（MNIST手写数字数据集） --------------------------
def load_data():
    # 加载训练集（60000张图，用于模型训练）
    train_dataset = datasets.MNIST(
        root='../dataset/mnist/',  # 数据集保存路径（上级目录的dataset文件夹）
        train=True,  # 标记为训练集
        download=True,  # 本地没有时自动从官网下载
        transform=train_transform  # 应用训练预处理
    )
    # 加载测试集（10000张图，用于评估模型性能）
    test_dataset = datasets.MNIST(
        root='../dataset/mnist/',
        train=False,  # 标记为测试集
        download=True,
        transform=train_transform  # 测试集用和训练集相同的预处理（除了不打乱）
    )
    # 创建数据加载器（批量加载+多线程）
    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True  # 训练集打乱顺序，避免模型学顺序特征
    )
    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False  # 测试集不打乱，结果可复现
    )
    print(f"Data loaded successfully: Train samples={len(train_dataset)}, Test samples={len(test_dataset)}")
    return train_loader, test_loader


# -------------------------- 4. 构建CNN模型（用于手写数字识别） --------------------------
class MNIST_CNN(torch.nn.Module):
    def __init__(self):
        super(MNIST_CNN, self).__init__()
        # 第一层卷积：输入1通道（灰度图）→输出10通道，卷积核5x5
        self.conv1 = torch.nn.Conv2d(in_channels=1, out_channels=10, kernel_size=5)
        # 第二层卷积：输入10通道→输出20通道，卷积核5x5
        self.conv2 = torch.nn.Conv2d(in_channels=10, out_channels=20, kernel_size=5)
        # 最大池化层：2x2窗口，缩小特征图尺寸（降维+减少计算量）
        self.pool = torch.nn.MaxPool2d(kernel_size=2)
        # 全连接层：输入320维（20通道×4×4特征图）→输出10维（对应10个数字类别）
        self.fc = torch.nn.Linear(in_features=320, out_features=10)

    # 前向传播（模型计算逻辑，必须实现）
    def forward(self, x):
        # x输入形状：[batch_size, 1, 28, 28]（批量数，通道数，图像高，图像宽）
        batch_size = x.size(0)  # 获取当前批量的样本数量
        
        # 第一层：卷积 → 池化 → ReLU激活（引入非线性）
        x = self.pool(F.relu(self.conv1(x)))  # 输出形状：[batch_size, 10, 12, 12]
        # 第二层：卷积 → 池化 → ReLU激活
        x = self.pool(F.relu(self.conv2(x)))  # 输出形状：[batch_size, 20, 4, 4]
        
        # 展平：把4维特征图（batch, channel, h, w）转2维向量（batch, channel×h×w）
        x = x.view(batch_size, -1)  # 输出形状：[batch_size, 320]（20×4×4=320）
        
        # 全连接层：输出10个类别的预测分数
        x = self.fc(x)  # 输出形状：[batch_size, 10]
        return x


# -------------------------- 5. 训练函数（单轮epoch训练逻辑） --------------------------
def train_one_epoch(model, train_loader, criterion, optimizer, epoch):
    model.train()  # 切换到训练模式（部分层如Dropout会生效）
    running_loss = 0.0  # 累计损失（用于打印进度）
    
    for batch_idx, (inputs, targets) in enumerate(train_loader, 0):
        # 把数据移到目标设备（GPU/CPU）
        inputs, targets = inputs.to(device), targets.to(device)
        
        # 1. 梯度清零（PyTorch默认累积梯度，每次迭代前需清空）
        optimizer.zero_grad()
        # 2. 前向传播：计算模型预测输出
        outputs = model(inputs)
        # 3. 计算损失（交叉熵损失，适合多分类）
        loss = criterion(outputs, targets)
        # 4. 反向传播：计算损失对参数的梯度
        loss.backward()
        # 5. 优化器更新参数（根据梯度调整权重）
        optimizer.step()
        
        # 累计损失并打印进度（每300个batch打印一次）
        running_loss += loss.item()  # .item()把Tensor转Python数值，避免内存占用
        if batch_idx % 300 == 299:  # 300个batch打印一次（索引从0开始，299是第300个）
            avg_loss = running_loss / 300  # 计算300个batch的平均损失
            print(f"[Epoch: {epoch+1}/{epochs}, Batch: {batch_idx+1}] Avg Loss: {avg_loss:.3f}")
            running_loss = 0.0  # 重置累计损失


# -------------------------- 6. 测试函数（评估准确率+返回结果用于绘图） --------------------------
def test_model(model, test_loader, criterion):
    model.eval()  # 切换到评估模式（禁用梯度计算，部分层如BatchNorm固定）
    total = 0  # 总样本数
    correct = 0  # 正确预测的样本数
    test_loss = 0.0  # 测试集总损失
    
    # 禁用梯度计算（测试时不需要求导，节省内存和时间）
    with torch.no_grad():
        for inputs, targets in test_loader:
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            
            # 计算测试损失
            loss = criterion(outputs, targets)
            test_loss += loss.item() * inputs.size(0)  # 乘以批量数，避免小批量偏差
            
            # 计算正确预测数：取输出最大值的索引（预测类别），和真实标签比较
            _, predicted = torch.max(outputs.data, dim=1)  # dim=1：沿类别维度取最大值
            total += targets.size(0)  # 累加当前批量的样本数
            correct += (predicted == targets).sum().item()  # 累加正确预测数
    
    # 计算测试集平均损失和准确率
    avg_test_loss = test_loss / total
    accuracy = 100 * correct / total
    print(f"\nTest Set Result: Avg Loss={avg_test_loss:.3f}, Accuracy={accuracy:.1f}% [{correct}/{total}]\n")
    return accuracy  # 返回准确率，用于绘制准确率曲线


# -------------------------- 7. 图像显示函数（显示初始图+识别结果） --------------------------
def show_images_with_results(model, test_loader):
    model.eval()  # 评估模式
    with torch.no_grad():
        # 取测试集第一个批量的数据（可修改为任意批量，这里用第一个）
        data_iter = iter(test_loader)  # 创建迭代器
        inputs, targets = next(data_iter)  # 取第一个批量（inputs: 图像，targets: 真实标签）
        inputs, targets = inputs.to(device), targets.to(device)
        
        # 模型预测
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, dim=1)  # 得到预测类别
        
        # 显示前10个样本（2行5列布局）
        for i in range(10):
            # 1. 处理图像：从Tensor转PIL格式（用于显示）
            img_tensor = inputs[i].cpu()  # 移到CPU（如果用GPU训练，必须移回CPU才能转PIL）
            img_tensor = inv_transform(img_tensor)  # 反归一化，还原为0-1的像素范围
            img = transforms.ToPILImage()(img_tensor)  # 转PIL图像（灰度图）
            
            # 2. 在子图中显示图像
            axs_img[i].imshow(img, cmap='gray')  # cmap='gray'：灰度显示，避免彩色偏差
            axs_img[i].axis('off')  # 关闭坐标轴，让图像显示更整洁
            
            # 3. 添加标签：真实值 vs 预测值（绿色=正确，红色=错误，直观区分）
            true_label = targets[i].item()  # 真实标签（0-9）
            pred_label = predicted[i].item()  # 预测标签（0-9）
            label_color = 'green' if true_label == pred_label else 'red'
            # 设置子图标题，换行显示真实值和预测值
            axs_img[i].set_title(f"True: {true_label}\nPred: {pred_label}", 
                                color=label_color, fontsize=10, pad=8)
        
        # 调整子图间距，避免标题重叠
        plt.tight_layout()
        # 设置总标题（在所有子图上方）
        fig_img.suptitle('MNIST Sample Images & Recognition Results', fontsize=14, y=1.02)
        # 刷新图像显示
        fig_img.canvas.draw()
        plt.pause(0.1)  # 暂停0.1秒，确保图像完全显示


# -------------------------- 8. 准确率曲线绘制函数（实时更新） --------------------------
def plot_accuracy_curve(epochs_list, accuracies_list):
    ax_acc.clear()  # 清空之前的曲线，避免重叠
    # 绘制准确率曲线：蓝色圆点+实线，线宽2，点大小6
    ax_acc.plot(epochs_list, accuracies_list, 'bo-', linewidth=2, markersize=6, label='Test Accuracy')
    # 设置坐标轴标签和标题
    ax_acc.set_xlabel('Training Epoch', fontsize=12)
    ax_acc.set_ylabel('Test Accuracy (%)', fontsize=12)
    ax_acc.set_title('Test Accuracy vs. Training Epochs', fontsize=14, pad=15)
    # 固定y轴范围（0-100%），让曲线变化更直观
    ax_acc.set_ylim(0, 100)
    # 添加网格线（虚线，透明度0.7），方便读取数值
    ax_acc.grid(True, linestyle='--', alpha=0.7)
    # 添加图例（显示曲线含义）
    ax_acc.legend(fontsize=10)
    # 刷新曲线显示
    plt.draw()
    plt.pause(0.1)


# -------------------------- 9. 主函数（整合全流程） --------------------------
def main():
    # 1. 加载数据
    train_loader, test_loader = load_data()
    
    # 2. 初始化模型、损失函数、优化器
    model = MNIST_CNN().to(device)  # 实例化模型并移到目标设备
    criterion = torch.nn.CrossEntropyLoss()  # 交叉熵损失（多分类常用）
    optimizer = optim.SGD(  # 随机梯度下降优化器
        model.parameters(),  # 要优化的模型参数（权重和偏置）
        lr=0.01,  # 学习率（控制参数更新步长）
        momentum=0.5  # 动量（加速收敛，减少震荡）
    )
    print(f"Model initialized successfully: {model}")
    
    # 3. 训练前：显示初始图像和随机预测结果（此时模型未训练，准确率低）
    print("\n=== Before Training: Random Predictions ===")
    show_images_with_results(model, test_loader)
    plt.pause(3)  # 暂停3秒，让你看清初始状态（红色标注多，因为预测随机）
    
    # 4. 训练循环（多轮epoch）
    epochs_list = []  # 存储已训练的epoch序号
    accuracies_list = []  # 存储对应epoch的测试准确率
    
    for epoch in range(epochs):
        print(f"\n=== Training Epoch {epoch+1}/{epochs} ===")
        # 单轮训练
        train_one_epoch(model, train_loader, criterion, optimizer, epoch)
        # 单轮测试，获取准确率
        current_acc = test_model(model, test_loader, criterion)
        # 记录epoch和准确率
        epochs_list.append(epoch + 1)
        accuracies_list.append(current_acc)
        # 实时更新准确率曲线
        plot_accuracy_curve(epochs_list, accuracies_list)
        # 显示当前epoch的图像识别结果（观察效果提升）
        show_images_with_results(model, test_loader)
    
    # 5. 训练结束：保持所有窗口显示（避免一闪而过）
    print("Training completed! Press 'Ctrl+C' or close windows to exit.")
    plt.ioff()  # 关闭交互模式
    plt.show()  # 保持图表显示


# -------------------------- 10. 运行主函数 --------------------------
if __name__ == '__main__':
    main()